program define poet , rclass
version 10.0

syntax varlist [if] [in]  , TREAT(string) [OTHER(varlist) full select(varlist)]
local vars  `varlist'
gettoken dv rhs : vars
qui {
marksample touse , novarlist
if "`full'"=="" {
markout `touse' `dv' `rhs' `treat' `other'
}
tempvar trresid lc

*** step 1: regress treatment on observables
regress `treat' `rhs' `other' if `touse'
predict `trresid' if e(sample) , residuals /*residuals from selection equation- calculated later because of e(sample) */
sum `dv' if e(sample)
scalar obs_var = `r(sd)'^2
*** step 2: regress outcome on rigorous residuals and observables 
reg `dv' `trresid' `rhs' `other' if `touse'

matrix beta=e(b)
scalar alpha = beta[1,1] /*rigorous ceofficient */
sum `treat' if e(sample)
scalar t_var = (r(sd))^2 /*variance in rigorous (var(r)) */
sum `trresid' if e(sample)
scalar t_tilde_var = (r(sd))^2 /*variance in selection residuals (var(ur)) */
scalar sim_adj = t_var/t_tilde_var
*** step 3: regress outcome on observables (i.e. the constrained equation, where alpha is constrained to equal zero)
reg `dv' `rhs' `other' if `touse'

scalar unob_var = e(rss)/(e(N)-1)  /*var(uc) */
*** Full linear combination
gen `lc' = 0 if e(sample)
foreach v of local rhs {
tempvar g_`v'
gen `g_`v'' = _b[`v']*`v' if e(sample)
replace `lc' = `lc' + _b[`v']*`v'  if e(sample)
 }

reg `lc' `treat' `other' if e(sample)
matrix beta8=e(b)
scalar g_full_r_ce = beta8[1,1]

reg `lc' `other' if e(sample)
scalar g_full_r_var = e(rss)/(e(N)-1) 
scalar ss_obs_full = g_full_r_ce/g_full_r_var
scalar unob_ce_full = ss_obs_full*unob_var
scalar sim_bias_full = sim_adj*unob_ce_full 
*return scalar altonji_ratio = alpha/sim_bias_full
scalar altonji_ratio = alpha/sim_bias_full
*return scalar imbens_ratio = ((alpha/sim_adj)^2)/(unob_ce_full^2)
scalar imbens_ratio = ((alpha/sim_adj)^2)/(unob_ce_full^2)
scalar r2_full = (t_var)*(g_full_r_ce^2)/(g_full_r_var)
scalar obs_select = altonji_ratio*unob_ce_full
*return scalar r2_full = r2_full
foreach v in obs_select g_full_r_var alpha unob_var g_full_r_ce unob_ce_full sim_adj /*ss_obs_full  sim_bias_full*/ altonji_ratio imbens_ratio /*r2_full*/ {
return scalar `v' = `v'
}
*** Individual components
foreach v of local rhs {

local rhs2 : list rhs - v
reg `g_`v'' `treat' `rhs2' `other' if e(sample)

scalar `v'_ce = _b[`treat']
reg `g_`v'' `rhs2' `other' if e(sample)
scalar `v'_var = e(rss)/(e(N)-1) 
scalar obs_`v' = `v'_ce/`v'_var
scalar unob_`v' = obs_`v'*unob_var
scalar sim_`v' = sim_adj*unob_`v'
scalar alt_`v' = alpha/sim_`v'
scalar imb_`v' = ((alpha/sim_adj)^2)/(unob_`v'^2)
scalar r2_`v' = (t_var)*(`v'_ce^2)/(`v'_var)

}

local o 25
local u 40
local c 54
local p 55
local a 70
local b 79
nois {
di as text " "
*di as text in red "Drumroll please ......"
*sleep 3000
di as text " "
di as text "Sensitivity analysis"
di as text "{hline 21}{c TT}{hline 64}"
di as text _col(22) "{c |}" _col(`o') %~10s "Obs." _col(`c') %~10s "Unobs. " _col(`a') %~17s "Implied Ratio" 
di as text   _col(22) "{c |}" _col(`o') "Cond. Mean" _col(`u') "Cond. Var" _col(`c') "Cond. Mean"  _col(`a') " Altonji" _col(`b') " Imbens"            "
di as text _col(22) "{c |}" _dup(64) "-"
}
/*
foreach v of local rhs {
nois di as text  "   " %17s abbrev("`v'",17)  _col(22) "{c |}" _col(`o') as result %7.0g `v'_ce _col(`u') as result %7.0g unob_`v' _col(`c') as result %7.0g `v'_var _col(`p') as result %8.0g r2_`v' _col(`a') as result %8.0g alt_`v'  _col(`b') as result %8.0g  imb_`v'  "
}
*/
nois di as text  _col(22) "{c |}" 
nois di as text   "   " %17s "Lin. combination"  _col(22) "{c |} " _col(`o') as result %7.0g g_full_r_ce _col(`u') as result %7.0g g_full_r_var _col(`c') as result %7.0g unob_ce_full  _col(`a') as result %8.0g altonji_ratio  _col(`b') as result %8.0g     imbens_ratio  "
nois di as text  _dup(21) "-" _col(22) "{c +}" _dup(64) "-"
nois di as text  "   " %17s "Alpha hat"  _col(22) "{c |} " _col(`o') as result %7.0g alpha "
nois di as text  "   " %17s "Var. disturb."  _col(22) "{c |} " _col(25) as result %7.0g unob_var "
nois di as text  "   " %17s "Adj. ratio"  _col(22) "{c |} " _col(`o') as result %7.0g `=1/sim_adj' "

nois di as text  "   " %17s "Selection on obs."  _col(22) "{c |} " _col(25) as result %7.0g obs_select "
*nois di as text  "   " %17s "Var. treat"  _col(22) "{c |} " _col(25) as result %7.0g t_var "
nois di as text "{hline 21}{c BT}{hline 64}"

}
end
